# 导入系统和文件操作相关库
import os  # 用于创建目录、路径操作
import json  # 用于保存和读取JSON格式的元数据
from datetime import datetime  # 用于记录训练时间戳

# 导入数值计算和可视化库
import numpy as np  # 数值计算库
import matplotlib  # 绘图库
matplotlib.use("Agg")  # 设置matplotlib后端为Agg，适合服务器环境（无GUI）
import matplotlib.pyplot as plt  # 绘图接口

# 导入机器学习相关库
from sklearn.datasets import fetch_california_housing  # 加载加州房价数据集
from sklearn.model_selection import train_test_split  # 划分训练测试集
from sklearn.ensemble import HistGradientBoostingRegressor  # 梯度提升回归器
from sklearn.pipeline import Pipeline  # 机器学习管道
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score  # 评估指标
import joblib  # 用于保存和加载模型


# 定义输出目录常量
MODELS_DIR = "models"  # 模型文件保存目录
ASSETS_DIR = "assets"  # 图表和资源文件保存目录

def ensure_dirs():
    """确保必要的目录存在，如果不存在则创建"""
    os.makedirs(MODELS_DIR, exist_ok=True)  # 创建模型目录，exist_ok=True表示目录已存在时不报错
    os.makedirs(ASSETS_DIR, exist_ok=True)  # 创建资源目录

def load_data():
    """
    加载加州房价数据集
    
    Returns:
        X: 特征数据（DataFrame）
        y: 目标变量（房价中位数）
        feature_names: 特征名称列表
        target_name: 目标变量名称
    """
    cali = fetch_california_housing(as_frame=True)  # 获取加州房价数据，as_frame=True返回DataFrame格式
    df = cali.frame.copy()  # 复制完整的数据框，包含特征和目标变量
    feature_names = cali.feature_names  # 获取特征名称列表：['MedInc', 'HouseAge', ...]
    target_name = cali.target_names[0]  # 获取目标变量名称：'MedHouseVal'
    X = df[feature_names]  # 提取特征数据
    y = df[target_name]  # 提取目标变量（房价中位数）
    return X, y, feature_names, target_name

def build_pipeline(random_state=42):
    """
    构建机器学习管道
    
    Args:
        random_state: 随机种子，确保结果可重现
        
    Returns:
        pipe: 包含预处理和模型的Pipeline对象
    """
    # 创建梯度提升回归器 - 强力树模型，免标准化，适合非线性关系
    model = HistGradientBoostingRegressor(
        learning_rate=0.1,      # 学习率：控制每棵树的贡献程度
        max_depth=None,         # 最大深度：None表示不限制深度
        max_iter=300,           # 最大迭代次数：即树的数量
        min_samples_leaf=20,    # 叶子节点最小样本数：防止过拟合
        l2_regularization=0.0,  # L2正则化：此处设为0，不使用正则化
        random_state=random_state  # 随机种子
    )
    # 创建管道：目前只包含模型，后续可添加预处理步骤
    pipe = Pipeline([
        ("model", model)  # 管道中的步骤：("步骤名称", 步骤对象)
    ])
    return pipe

def plot_pred_vs_true(y_true, y_pred, path):
    """
    绘制预测值vs真实值散点图
    
    Args:
        y_true: 真实值数组
        y_pred: 预测值数组
        path: 图片保存路径
    """
    plt.figure(figsize=(6, 6))  # 创建6x6英寸的正方形图表
    # 绘制散点图：真实值为x轴，预测值为y轴
    plt.scatter(y_true, y_pred, s=8, alpha=0.5, edgecolors="none")
    # 计算数据范围，用于绘制理想预测线（y=x）
    lims = [min(y_true.min(), y_pred.min()), max(y_true.max(), y_pred.max())]
    plt.plot(lims, lims, "r--", linewidth=1)  # 绘制红色虚线，表示完美预测线
    plt.xlabel("True Median House Value")     # x轴标签
    plt.ylabel("Predicted Median House Value") # y轴标签
    plt.title("Predicted vs. True")           # 图表标题
    plt.tight_layout()  # 自动调整布局，防止标签重叠
    plt.savefig(path, dpi=150)  # 保存图片，dpi=150设置较高分辨率
    plt.close()  # 关闭图表，释放内存

def plot_residuals(y_true, y_pred, path):
    """
    绘制残差分布直方图
    
    Args:
        y_true: 真实值数组
        y_pred: 预测值数组
        path: 图片保存路径
    """
    residuals = y_true - y_pred  # 计算残差：真实值 - 预测值
    plt.figure(figsize=(6, 4))   # 创建6x4英寸的图表
    # 绘制残差直方图：50个区间，蓝色，半透明
    plt.hist(residuals, bins=50, color="#4C72B0", alpha=0.85)
    plt.xlabel("Residuals (True - Pred)")  # x轴标签：残差值
    plt.ylabel("Count")                    # y轴标签：频次
    plt.title("Residuals Distribution")    # 图表标题
    plt.tight_layout()  # 自动调整布局
    plt.savefig(path, dpi=150)  # 保存图片
    plt.close()  # 关闭图表

def plot_feature_importance(model, feature_names, path):
    """
    绘制特征重要性条形图
    
    Args:
        model: 训练好的模型对象
        feature_names: 特征名称列表
        path: 图片保存路径
    """
    # 检查模型是否有feature_importances_属性（HistGradientBoostingRegressor有此属性）
    if hasattr(model, "feature_importances_"):
        importances = model.feature_importances_  # 获取特征重要性数组
        idx = np.argsort(importances)  # 按重要性排序，返回索引数组（从小到大）
        plt.figure(figsize=(7, 5))     # 创建7x5英寸的图表
        # 绘制水平条形图：按重要性从低到高排列
        plt.barh(np.array(feature_names)[idx], importances[idx], color="#55A868")
        plt.xlabel("Importance")          # x轴标签：重要性
        plt.title("Feature Importances")  # 图表标题
        plt.tight_layout()  # 自动调整布局
        plt.savefig(path, dpi=150)  # 保存图片
        plt.close()  # 关闭图表

def main():
    """主函数：执行完整的模型训练流程"""
    ensure_dirs()  # 确保输出目录存在
    X, y, feature_names, target_name = load_data()  # 加载数据集

    # 划分训练集和测试集：80%训练，20%测试，随机种子42确保可重现
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    pipe = build_pipeline()  # 构建机器学习管道
    pipe.fit(X_train, y_train)  # 在训练集上训练模型
    y_pred = pipe.predict(X_test)  # 在测试集上进行预测

    # 计算评估指标
    rmse = float(np.sqrt(mean_squared_error(y_test, y_pred)))  # 均方根误差
    mae = float(mean_absolute_error(y_test, y_pred))          # 平均绝对误差
    r2 = float(r2_score(y_test, y_pred))                      # R²决定系数

    # 创建包含训练信息和评估结果的元数据字典
    metrics = {
        "timestamp": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ"),  # UTC时间戳
        "n_train": int(X_train.shape[0]),    # 训练样本数量
        "n_test": int(X_test.shape[0]),      # 测试样本数量
        "features": feature_names,           # 特征名称列表
        "target": target_name,               # 目标变量名称
        "metrics": {                         # 模型性能指标
            "RMSE": rmse,
            "MAE": mae,
            "R2": r2
        },
        "model": {                           # 模型信息
            "type": "HistGradientBoostingRegressor",
            "sklearn_version": __import__("sklearn").__version__  # sklearn版本
        }
    }

    # 保存模型与元数据到文件
    model_path = os.path.join(MODELS_DIR, "california_housing_pipeline.joblib")
    joblib.dump(pipe, model_path)  # 保存训练好的管道对象
    
    # 保存特征名称列表（用于后续预测时的特征对齐）
    with open(os.path.join(MODELS_DIR, "feature_names.json"), "w", encoding="utf-8") as f:
        json.dump(feature_names, f, ensure_ascii=False, indent=2)
    
    # 保存训练指标和元数据
    with open(os.path.join(MODELS_DIR, "metrics.json"), "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)

    # 保存示例输入数据（用于API测试）
    sample = X_test.iloc[0].to_dict()  # 将测试集第一个样本转为字典格式
    with open(os.path.join(MODELS_DIR, "sample_input.json"), "w", encoding="utf-8") as f:
        json.dump(sample, f, ensure_ascii=False, indent=2)

    # 生成可视化图表
    # 1. 预测值vs真实值散点图
    plot_pred_vs_true(
        y_test.to_numpy(),  # 将pandas Series转为numpy数组
        y_pred,
        os.path.join(ASSETS_DIR, "pred_vs_true.png")
    )
    
    # 2. 残差分布直方图
    plot_residuals(
        y_test.to_numpy(),
        y_pred,
        os.path.join(ASSETS_DIR, "residuals.png")
    )
    
    # 3. 特征重要性图表
    base_model = pipe.named_steps["model"]  # 从管道中提取底层模型对象
    plot_feature_importance(
        base_model, feature_names, os.path.join(ASSETS_DIR, "feature_importance.png")
    )

    # 输出训练结果摘要
    print("Training complete.")  # 训练完成提示
    print(f"RMSE: {rmse:.4f}, MAE: {mae:.4f}, R2: {r2:.4f}")  # 打印评估指标
    print(f"Saved model to: {model_path}")  # 打印模型保存路径
    print("Plots saved in assets/")  # 打印图表保存位置

if __name__ == "__main__":
    main()  